#include <opencv2/opencv.hpp>
#include <iostream>
#include <filesystem>

namespace fs = std::filesystem;

std::vector<std::string> get_files_with_extension(const std::string& directory, const std::string& extension) {
    std::vector<std::string> files;
    for (const auto& entry : fs::directory_iterator(directory)) {
        if (entry.is_regular_file() && entry.path().extension() == extension) {
            files.push_back(entry.path().filename().string());
        }
    }
    return files;
}

int detect_license_plate(const std::string& imgPath) {
    // 读取图像
    cv::Mat img = cv::imread(imgPath);    
    if (img.empty()) {
        std::cerr << "无法打开图像文件!" << std::endl;
        return -1;
    }

    // 将图像转换为HSV颜色空间
    cv::Mat hsv;
    cv::cvtColor(img, hsv, cv::COLOR_BGR2HSV);

    // 设置蓝色的HSV范围
    cv::Scalar lowerBlue(100, 150, 50);  // 蓝色的低阈值
    cv::Scalar upperBlue(140, 255, 255); // 蓝色的高阈值

    // 创建掩膜，提取蓝色部分
    cv::Mat mask;
    cv::inRange(hsv, lowerBlue, upperBlue, mask);

    // 对掩膜进行形态学操作，去除噪声
    cv::Mat morphMask;
    cv::morphologyEx(mask, morphMask, cv::MORPH_CLOSE, cv::Mat());

    // 使用掩膜从原图中提取蓝色区域
    cv::Mat blueRegion;
    cv::bitwise_and(img, img, blueRegion, morphMask);

    // 转换为灰度图并进行边缘检测
    cv::Mat gray;
    cv::cvtColor(blueRegion, gray, cv::COLOR_BGR2GRAY);
    cv::Mat edges;
    cv::Canny(gray, edges, 100, 200);

    // 寻找轮廓
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(edges, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    // 遍历轮廓，找出车牌的矩形区域
    // 判断框选区域是否大小合适
    int imgHeight = img.rows;
    int imgWidth = img.cols;
    for (size_t i = 0; i < contours.size(); ++i) {
        cv::Rect boundingBox = cv::boundingRect(contours[i]);

        // 设定车牌区域的宽高比例来进行筛选
        float aspectRatio = (float)boundingBox.width / boundingBox.height;
        if (aspectRatio > 2.0 && aspectRatio < 5.0 && (boundingBox.width > 0.05 * imgWidth && boundingBox.height > 0.05* imgHeight)) { // 假设车牌的宽高比在这个范围内
            // 在图像上绘制矩形框
            cv::rectangle(img, boundingBox, cv::Scalar(0, 255, 0), 2);
        }
    }


    // 显示结果
    fs::path p(imgPath);
    std::string outputFilename = "dst_" + (std::string)p.filename();
    cv::imwrite(outputFilename, img);
    return 0;
}

int save_image(const std::string& imgPath, const cv::Mat& img) {
    fs::path p(imgPath);
    std::string outputFilename = "dst_" + (std::string)p.filename();
    cv::imwrite(outputFilename, img);
    return 0;
}

int detect_license_plate_affine_transformation(const std::string& imgPath) {
    // 读取图像
    cv::Mat img = cv::imread(imgPath);
    if (img.empty()) {
        std::cerr << "无法打开图像文件!" << std::endl;
        return -1;
    }

    // 将图像转换为HSV颜色空间
    cv::Mat hsv;
    cv::cvtColor(img, hsv, cv::COLOR_BGR2HSV);

    // 设置蓝色的HSV范围
    cv::Scalar lowerBlue(100, 150, 50);  // 蓝色的低阈值
    cv::Scalar upperBlue(140, 255, 255); // 蓝色的高阈值

    // 创建掩膜，提取蓝色部分
    cv::Mat mask;
    cv::inRange(hsv, lowerBlue, upperBlue, mask);

    // 对掩膜进行形态学操作，去除噪声
    cv::Mat morphMask;
    cv::morphologyEx(mask, morphMask, cv::MORPH_CLOSE, cv::Mat());

    // 使用掩膜从原图中提取蓝色区域
    cv::Mat blueRegion;
    cv::bitwise_and(img, img, blueRegion, morphMask);

    // 转换为灰度图并进行边缘检测
    cv::Mat gray;
    cv::cvtColor(blueRegion, gray, cv::COLOR_BGR2GRAY);
    cv::Mat edges;
    cv::Canny(gray, edges, 100, 200);

    // 寻找轮廓
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(edges, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    // 遍历轮廓，找出车牌的矩形区域
    int imgHeight = img.rows;
    int imgWidth = img.cols;
    for (size_t i = 0; i < contours.size(); ++i) {
        cv::Rect boundingBox = cv::boundingRect(contours[i]);

        // 设定车牌区域的宽高比例来进行筛选
        float aspectRatio = (float)boundingBox.width / boundingBox.height;
        if (aspectRatio > 2.0 && aspectRatio < 5.0) { // 假设车牌的宽高比在这个范围内
            // 获取矩形的四个角点
            cv::Point2f rect_points[4];
            cv::RotatedRect rotatedRect = cv::minAreaRect(contours[i]);
            rotatedRect.points(rect_points);

            // 对车牌进行透视变换，将其转换为矩形
            std::vector<cv::Point2f> srcPoints = {rect_points[0], rect_points[1], rect_points[2], rect_points[3]};
            std::vector<cv::Point2f> dstPoints = {
                cv::Point2f(0, 0), 
                cv::Point2f(boundingBox.width - 1, 0), 
                cv::Point2f(boundingBox.width - 1, boundingBox.height - 1), 
                cv::Point2f(0, boundingBox.height - 1)
            };

            // 计算透视变换矩阵
            cv::Mat perspectiveMatrix = cv::getPerspectiveTransform(srcPoints, dstPoints);

            // 进行透视变换
            cv::Mat warpedPlate;
            cv::warpPerspective(img, warpedPlate, perspectiveMatrix, boundingBox.size());

            // 显示透视校正后的车牌
            fs::path p(imgPath);
            std::string outputFilename = "dst_" + (std::string)p.filename();
            cv::imwrite(outputFilename, img);
        }
    }

    return 0;
}

int main() {
    std::string directory = "./assets/";
    std::string extension = ".jpeg";

    auto files = get_files_with_extension(directory, extension);
    std::vector<std::string> fullPaths;
    for(const auto& file : files) {
        std::string fullPath = directory + file;
        fullPaths.push_back(fullPath);
        std::cout << fullPath << std::endl;
        detect_license_plate(fullPath);
    }
    
    return 0;
}